import os
import sys
import json
import time
import re
import urllib3
import hashlib
import yaml
import requests


# 禁用 InsecureRequestWarning
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

class ConsulNodeRegistration:
    def __init__(self, ip_file_path="ip.txt", consul_url="http://consul.opsx.top"):
        self.ip_file_path = ip_file_path
        self.consul_url = consul_url
        
        # 检查IP文件是否存在
        if not os.path.isfile(self.ip_file_path):
            print(f"错误: IP文件 '{self.ip_file_path}' 不存在")
            sys.exit(1)
            
        print(f"初始化成功, 将从 {self.ip_file_path} 读取IP地址列表")
        

    def read_ip_file(self):
        """从文件中读取IP地址列表"""
        try:
            with open(self.ip_file_path, 'r') as f:
                # 读取所有行并移除空白字符
                ip_list = [line.strip() for line in f.readlines() if line.strip()]
            return ip_list
        except Exception as e:
            print(f"读取IP文件时出错: {str(e)}")
            return []
    
  
    def register_node_exporters(self):
        """从IP文件读取IP地址并注册到Consul"""
        # 读取IP列表
        ip_list = self.read_ip_file()
        
        if not ip_list:
            print("没有找到可用的IP地址")
            return
            
        print(f"找到 {len(ip_list)} 个IP地址，开始注册到Consul")
        
        for index, ip in enumerate(ip_list):
            try:
                # 生成唯一的节点名称
                node_name = f"node-{index+1}"
                
                print(f"正在注册: IP: {ip} \t 节点: {node_name}")
                
                # 定义服务注册的数据
                data = {
                    "id": f"{node_name}-exporter",
                    "name": "node-exporter",
                    "address": ip,
                    "port": 9100,
                    "checks": [{
                        "http": f"http://{ip}:9100/metrics",
                        "interval": "5s"
                    }]
                }
                
                # 发送 PUT 请求以注册服务
                response = requests.put(f"{self.consul_url}/v1/agent/service/register", json=data)

                # 检查响应状态
                if response.status_code == 200:
                    print(f"服务 {node_name} 注册成功.")
                else:
                    print(f"无法注册服务 {node_name}. 状态码: {response.status_code}")
                    print(response.text)
            except Exception as e:
                print(f"注册IP {ip} 时出错: {str(e)}")

    def clean_failed_instances(self):
        """清理失效的服务实例"""
        time.sleep(3)
        try:
            response = requests.get(f"{self.consul_url}/v1/health/state/critical")
            if response.status_code == 200:
                instances = response.json()
                for instance in instances:
                    if instance['Status'] == 'critical':  # 如果实例状态为严重
                        service_id = instance['ServiceID']
                        requests.put(f"{self.consul_url}/v1/agent/service/deregister/{service_id}")
                        print(f"已清理失效实例ID: {service_id}")
            else:
                print(f"无法从 Consul API 获取数据。状态码：{response.status_code}")
        except Exception as e:
            print(f"清理失效实例时出错: {str(e)}")
                


if __name__ == "__main__":
    # 默认从当前目录下的ip.txt文件读取IP地址
    # 可以通过命令行参数指定不同的文件路径
    ip_file = "ip.txt"
    if len(sys.argv) > 1:
        ip_file = sys.argv[1]
    
    consul = ConsulNodeRegistration(ip_file_path=ip_file)
    consul.register_node_exporters()
    consul.clean_failed_instances()

